{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2D/3D acoustic problem\n",
    "\n",
    "[![Download Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_notebook.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/master/mindflow/zh_cn/cfd_solver/mindspore_acoustic.ipynb)&emsp;[![Download Sample Codes](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_download_code.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/master/mindflow/zh_cn/cfd_solver/mindspore_acoustic.py)&emsp;[![View Source Files](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_source.png)](https://gitee.com/mindspore/docs/blob/master/docs/mindflow/docs/source_zh_cn/cfd_solver/acoustic.ipynb)\n",
    "\n",
    "## Environment Installation\n",
    "\n",
    "This case requires **MindSpore >= 2.4.0** version. Please refer to [MindSpore Installation](https://www.mindspore.cn/install) for details.\n",
    "\n",
    "In addition, you need to install **MindFlow >=0.4.0** version. If it is not installed in the current environment, please follow the instructions below to choose the backend and version for installation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mindflow_version = \"0.4.0\"  # update if needed\n",
    "\n",
    "# Only NPU is supported.\n",
    "!pip uninstall -y mindflow-ascend\n",
    "!pip install mindflow-ascend==$mindflow_version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "Solving the acoustic wave equation is a core technology in fields such as medical ultrasound and geological exploration. Large-scale acoustic wave equation solvers face challenges in terms of computational power and storage. Solvers for the wave equation generally use either frequency domain algorithms or time domain algorithms. The representative time domain algorithm is the Time Domain Finite Difference (TDFD) method, while frequency domain algorithms include Frequency Domain Finite Difference (FDFD), Finite Element Method (FEM), and Convergent Born Series (CBS) iterative method. The CBS method, due to its low memory requirement and absence of dispersion error, has gained widespread attention in engineering and academia. In particular, [Osnabrugge et al. (2016)](https://linkinghub.elsevier.com/retrieve/pii/S0021999116302595) have addressed the convergence issue of this method, expanding the application prospects of the CBS method.\n",
    "\n",
    "This case study will demonstrate how to invoke the CBS API provided by MindFlow to solve the two/three-dimensional acoustic wave equation."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Description\n",
    "\n",
    "In solving the acoustic wave equation, the input parameters are the velocity field and source information, and the output is the spatiotemporal distribution of the wave field.\n",
    "\n",
    "The expression of the two/three-dimensional acoustic wave equation is as follows\n",
    "\n",
    "| Time Domain Expression                                            | Frequency Domain Expression                                        |\n",
    "| ----------------------------------------------------- | ------------------------------------------------- |\n",
    "| $\\frac{\\partial^2u}{\\partial t^2} - c^2 \\Delta u = f$ | $-\\omega^2 \\hat{u} - c^2 \\Delta\\hat{u} = \\hat{f}$ |\n",
    "\n",
    "Where\n",
    "\n",
    "- $u(\\bold{x},t) \\;\\; [L]$ is the deformation displacement (pressure divided by density), a scalar.\n",
    "- $c(\\bold{x}) \\;\\; [L/T]$ is the wave velocity, a scalar.\n",
    "- $f(\\bold{x},t) \\;\\; [L/T^2]$ is the excitation source (volume distributed force), a scalar.\n",
    "\n",
    "In practical solving, in order to reduce the parameter dimension, the parameters are generally made dimensionless first, and then the dimensionless equations and parameters are solved, and finally the dimensional solutions are restored. By selecting $\\omega$, $\\hat{f}$, and $d$ (grid spacing, which requires equal spacing in all directions) to nondimensionalize the frequency domain equation, we can obtain the dimensionless frequency domain equation:\n",
    "\n",
    "$$\n",
    "u^* + c^{*2} \\tilde{\\Delta} + f^* = 0\n",
    "$$\n",
    "\n",
    "Where\n",
    "\n",
    "- $u^* = \\hat{u} \\omega^2 / \\hat{f}$ is the dimensionless deformation displacement.\n",
    "- $c^* = c / (\\omega d)$ is the dimensionless wave velocity.\n",
    "- $\\tilde{\\Delta}$ is the normalized Laplace operator, which is the Laplace operator when the grid spacing is 1.\n",
    "- $f^*$ the mask that marks the source position, with a value of 1 at the source and 0 at other positions.\n",
    "\n",
    "The `src` package in this case can be downloaded at [src](https://gitee.com/mindspore/mindscience/tree/master/MindFlow/applications/cfd/acoustic/src)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import mindspore as ms\n",
    "from mindspore import Tensor\n",
    "\n",
    "from mindflow.utils import load_yaml_config\n",
    "\n",
    "from cbs.cbs import CBS\n",
    "from src import visual\n",
    "from solve_acoustic import solve_cbs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define input parameters and output sampling method\n",
    "\n",
    "The required inputs for this case are dimensional 2D/3D velocity field, source location list, and source waveform. The input file name is specified in the [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) (2D) or [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) (3D) file. For user convenience, pre-set inputs are provided [here](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/cfd/acoustic). Please download the data and put them in `./dataset` in the case directory. The data include the velocity field `velocity_2d.npy`, source location list `srclocs_2d.csv`, and source waveform `srcwaves_2d.csv` in 2D, and the velocity field `velocity_3d.npy`, source location list `srclocs_3d.csv`, and source waveform `srcwaves_3d.csv` in 3D. Users can modify the input parameters based on the input file format.\n",
    "\n",
    "The output is a spatiotemporal distribution of the wavefield. To specify how the output is sampled in time and frequency, parameters such as `dt` and `nt` need to be specified in the [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) or [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) file.\n",
    "\n",
    "Since the sampling rate of the input source waveform in time may differ from the required output, interpolation needs to be performed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ms.set_context(device_target='Ascend', device_id=0, mode=ms.GRAPH_MODE)\n",
    "\n",
    "# Space dimension\n",
    "dim = 2\n",
    "\n",
    "if dim == 2:\n",
    "    config = load_yaml_config(\"./config_2d.yaml\")\n",
    "else:\n",
    "    config = load_yaml_config(\"./config_3d.yaml\")\n",
    "\n",
    "data_config = config['data']\n",
    "solve_config = config['solve']\n",
    "summary_config = config['summary']\n",
    "\n",
    "# Coarsen rate for reducing scale\n",
    "rate = solve_config['coarsen_rate']\n",
    "\n",
    "# Read time & frequency points\n",
    "dt = solve_config['dt']\n",
    "nt = solve_config['nt']\n",
    "dt *= rate # Increase the time iteraion step size\n",
    "nt //= rate\n",
    "receivers = None\n",
    "\n",
    "# Read velocity array\n",
    "velo = np.load(os.path.join(data_config['root_dir'], data_config['velocity_field']))\n",
    "assert dim == len(velo.shape), \"dim and dimension of velocity should be equal\"\n",
    "dx = data_config['velocity_dx']\n",
    "dz = data_config['velocity_dz']\n",
    "if dim == 2:\n",
    "    dxs = (rate*dz, rate*dx)\n",
    "    velo = velo[::rate, ::rate] # Coarsen velocity field\n",
    "else:\n",
    "    dy = data_config['velocity_dy']\n",
    "    dxs = (rate*dz, rate*dy, rate*dx)\n",
    "    velo = velo[::rate, ::rate, ::rate] # Coarsen velocity field\n",
    "\n",
    "# Read source locations\n",
    "df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_locations']), index_col=0)\n",
    "if dim == 2:\n",
    "    slocs = df[['y', 'x']].values # Shape (ns, 2)\n",
    "else:\n",
    "    slocs = df[['z', 'y', 'x']].values # Shape (ns, 3)\n",
    "\n",
    "# Read & interp source wave\n",
    "df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_wave']))\n",
    "inter_func = interp1d(df.t, df.f, kind='cubic', bounds_error=False, fill_value=0)\n",
    "\n",
    "ts = np.arange(nt) * dt\n",
    "omegas_all = np.fft.rfftfreq(nt) * (2 * np.pi / dt)\n",
    "\n",
    "# Interpolation source wave\n",
    "src_waves = inter_func(ts) # Shape (nt)\n",
    "src_amplitudes = np.fft.rfft(src_waves) # Shape (nt//2+1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Select desired frequency points\n",
    "\n",
    "With the output sampling method determined, all the desired frequency points are in turn determined. However, in order to reduce computational load, it is also possible to select only a portion of the frequency points for calculation, while obtaining the remaining frequency points through interpolation. The specific frequency point downsampling method is specified by the `downsample_mode` and `downsample_rate` in the [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) (2D) or [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) (3D) file. The default is no downsampling, which means solving all frequency points except $\\omega=0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select omegas\n",
    "no = len(omegas_all) // solve_config['downsample_rate']\n",
    "\n",
    "if solve_config['downsample_mode'] == 'exp':\n",
    "    omegas_sel = np.exp(np.linspace(np.log(omegas_all[1]), np.log(omegas_all[-1]), no))\n",
    "elif solve_config['downsample_mode'] == 'square':\n",
    "    omegas_sel = np.linspace(omegas_all[1]**.5, omegas_all[-1]**.5, no)**2\n",
    "else:\n",
    "    omegas_sel = np.linspace(omegas_all[1], omegas_all[-1], no)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform simulation\n",
    "\n",
    "Define the relevant arrays as Tensors, call `solve_cbs()`, and execute the solution on the NPU. Due to memory limitations, the solution process is executed in batches in the frequency domain. The number of batches is specified by the user in `config_2d.yaml` or `config_3d.yaml` and does not need to be divisible by the number of frequency points (allowing the size of the last batch to be different from the other batches). The thickness and type of absorbing boundaries in each direction can be specified via parameters `pml_size` and `btype` in the file [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) or [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml). After the solution is completed, the frequency domain solution results will be saved to the file `u_star.npy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Send to NPU and perform computation\n",
    "os.makedirs(summary_config['root_dir'], exist_ok=True)\n",
    "velo = Tensor(velo, dtype=ms.float32, const_arg=True)\n",
    "\n",
    "dxs_nd = tuple(d / dxs[-1] for d in dxs) # Nondimensional dxs\n",
    "pml_size = tuple(solve_config['pml_size'])\n",
    "btype = solve_config['btype']\n",
    "cbs = CBS(velo.shape, dxs=dxs_nd, pml_size=pml_size, alpha=1., rampup=12,\n",
    "          btype=btype, remove_pml=False)\n",
    "\n",
    "# Shape (ns, no, len(receiver_zs), nx) or (ns, no, len(receiver_zs), ny, nx)\n",
    "ur, ui, errs = solve_cbs(\n",
    "    cbs, velo, slocs, omegas_sel, receivers=receivers, dxs=dxs, n_batches=solve_config['n_batches'])\n",
    "\n",
    "u_star = ur.numpy() + 1j * ui.numpy() # Shape (ns, no, len(krs), nx) or (ns, no, len(krs), ny, nx)\n",
    "\n",
    "np.save(os.path.join(summary_config['root_dir'], 'u_star.npy'), np.squeeze(u_star))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Post-processing\n",
    "\n",
    "CBS solves the dimensionless frequency domain equation, but downstream tasks often require observing the evolution process of dimensional wavefields in the time domain. Therefore, the final solution is restored to dimensional and converted back to the time domain. The restoration method is given by $\\hat{u} = u^* \\hat{f} / \\omega^2$. If downsampling is performed on the frequency points in the \"Select desired frequency points\" step, interpolation along the frequency direction is required here to restore the solutions for all frequency points. Then, perform a Fourier inverse transform on the dimensional frequency domain wavefield $\\hat{u}$ to obtain the time domain wavefield $u$. Save the time domain wavefield to the file `u_time.npy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Recover dimension and interpolate to full frequency domain\n",
    "ones = (1,) * dim\n",
    "u_star /= omegas_sel.reshape(-1, *ones)**2\n",
    "u_star = interp1d(omegas_sel, u_star, axis=1, kind='cubic', bounds_error=False, fill_value=0)(omegas_all)\n",
    "u_star *= src_amplitudes.reshape(-1, *ones)\n",
    "\n",
    "# Transform to time domain\n",
    "u_time = np.fft.irfft(u_star, axis=1)\n",
    "np.save(os.path.join(summary_config['root_dir'], 'u_time.npy'), u_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, read the time-domain wave field and visualize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the result\n",
    "u_time = np.load(os.path.join(summary_config['root_dir'], 'u_time.npy'))\n",
    "if dim == 2:\n",
    "    visual.anim(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif'))\n",
    "else:\n",
    "    visual.anim3d(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif'))\n",
    "visual.plot_errs(errs, os.path.join(summary_config['root_dir'], 'errors.png'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6 (main, Nov 14 2022, 16:10:14) [GCC 11.3.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
